# encoding: utf-8

import numpy as np
import struct
import matplotlib.pyplot as plt


class LoadHandWriteData:
    # 训练集文件
    train_images_idx3_ubyte_file = 'data/train-images.idx3-ubyte'
    # 训练集标签文件
    train_labels_idx1_ubyte_file = 'data/train-labels.idx1-ubyte'

    # 测试集文件
    test_images_idx3_ubyte_file = 'data/t10k-images.idx3-ubyte'
    # 测试集标签文件
    test_labels_idx1_ubyte_file = 'data/t10k-labels.idx1-ubyte'

    @staticmethod
    def decode_idx3_ubyte(idx3_ubyte_file):
        """
        解析idx3文件的通用函数
        :param idx3_ubyte_file: idx3文件路径
        :return: 数据集
        """
        # 读取二进制数据
        bin_data = open(idx3_ubyte_file, 'rb').read()

        # 解析文件头信息，依次为魔数、图片数量、每张图片高、每张图片宽
        offset = 0
        fmt_header = '>iiii'
        magic_number, num_images, num_rows, num_cols = struct.unpack_from(fmt_header, bin_data, offset)

        # 解析数据集
        image_size = num_rows * num_cols
        offset += struct.calcsize(fmt_header)
        fmt_image = '>' + str(image_size) + 'B'
        images = np.empty((num_images, num_rows, num_cols))
        for i in range(num_images):
            images[i] = np.array(struct.unpack_from(fmt_image, bin_data, offset)).reshape((num_rows, num_cols))
            offset += struct.calcsize(fmt_image)
        return images

    @staticmethod
    def decode_idx1_ubyte(idx1_ubyte_file):
        """
        解析idx1文件的通用函数
        :param idx1_ubyte_file: idx1文件路径
        :return: 数据集
        """
        # 读取二进制数据
        bin_data = open(idx1_ubyte_file, 'rb').read()

        # 解析文件头信息，依次为魔数和标签数
        offset = 0
        fmt_header = '>ii'
        magic_number, num_images = struct.unpack_from(fmt_header, bin_data, offset)

        # 解析数据集
        offset += struct.calcsize(fmt_header)
        fmt_image = '>B'
        labels = np.empty(num_images)
        for i in range(num_images):
            labels[i] = struct.unpack_from(fmt_image, bin_data, offset)[0]
            offset += struct.calcsize(fmt_image)
        return labels

    @staticmethod
    def load_train_images(idx_ubyte_file=train_images_idx3_ubyte_file):
        """
        TRAINING SET IMAGE FILE (train-images-idx3-ubyte):
        [offset] [type]          [value]          [description]
        0000     32 bit integer  0x00000803(2051) magic number
        0004     32 bit integer  60000            number of images
        0008     32 bit integer  28               number of rows
        0012     32 bit integer  28               number of columns
        0016     unsigned byte   ??               pixel
        0017     unsigned byte   ??               pixel
        ........
        xxxx     unsigned byte   ??               pixel
        Pixels are organized row-wise. Pixel values are 0 to 255. 0 means background (white),
        255 means foreground (black).

        :param idx_ubyte_file: idx文件路径
        :return: n*row*col维np.array对象，n为图片数量
        """
        return LoadHandWriteData.decode_idx3_ubyte(idx_ubyte_file)

    @staticmethod
    def load_train_labels(idx_ubyte_file=train_labels_idx1_ubyte_file):
        """
        TRAINING SET LABEL FILE (train-labels-idx1-ubyte):
        [offset] [type]          [value]          [description]
        0000     32 bit integer  0x00000801(2049) magic number (MSB first)
        0004     32 bit integer  60000            number of items
        0008     unsigned byte   ??               label
        0009     unsigned byte   ??               label
        ........
        xxxx     unsigned byte   ??               label
        The labels values are 0 to 9.

        :param idx_ubyte_file: idx文件路径
        :return: n*1维np.array对象，n为图片数量
        """
        return LoadHandWriteData.decode_idx1_ubyte(idx_ubyte_file)

    @staticmethod
    def load_test_images(idx_ubyte_file=test_images_idx3_ubyte_file):
        """
        TEST SET IMAGE FILE (t10k-images-idx3-ubyte):
        [offset] [type]          [value]          [description]
        0000     32 bit integer  0x00000803(2051) magic number
        0004     32 bit integer  10000            number of images
        0008     32 bit integer  28               number of rows
        0012     32 bit integer  28               number of columns
        0016     unsigned byte   ??               pixel
        0017     unsigned byte   ??               pixel
        ........
        xxxx     unsigned byte   ??               pixel
        Pixels are organized row-wise. Pixel values are 0 to 255. 0 means background (white),
        255 means foreground (black).

        :param idx_ubyte_file: idx文件路径
        :return: n*row*col维np.array对象，n为图片数量
        """
        return LoadHandWriteData.decode_idx3_ubyte(idx_ubyte_file)

    @staticmethod
    def load_test_labels(idx_ubyte_file=test_labels_idx1_ubyte_file):
        """
        TEST SET LABEL FILE (t10k-labels-idx1-ubyte):
        [offset] [type]          [value]          [description]
        0000     32 bit integer  0x00000801(2049) magic number (MSB first)
        0004     32 bit integer  10000            number of items
        0008     unsigned byte   ??               label
        0009     unsigned byte   ??               label
        ........
        xxxx     unsigned byte   ??               label
        The labels values are 0 to 9.

        :param idx_ubyte_file: idx文件路径
        :return: n*1维np.array对象，n为图片数量
        """
        return LoadHandWriteData.decode_idx1_ubyte(idx_ubyte_file)

    @staticmethod
    def run():
        train_images = LoadHandWriteData.load_train_images()
        train_labels = LoadHandWriteData.load_train_labels()
        # test_images = load_test_images()
        # test_labels = load_test_labels()

        # 查看前十个数据及其标签以读取是否正确
        for i in range(10):
            print(train_labels[i])
            plt.imshow(train_images[i], cmap='gray')
            print(train_images[i])
            plt.show()
        print('done')

    @staticmethod
    def load_data():
        x_train = LoadHandWriteData.load_train_images()
        y_train = LoadHandWriteData.load_train_labels()
        x_test = LoadHandWriteData.load_test_images()
        y_test = LoadHandWriteData.load_test_labels()
        return (x_train, y_train), (x_test, y_test)


if __name__ == '__main__':
    LoadHandWriteData.run()
